#include "evdeploy/runtime/sophon/sophon_common.h"
#include "evdeploy/utils/report/tool_report.h"
#include <chrono>
#include <numeric>

namespace ev
{

    SophonModel::SophonModel() { EVLOG(INFO) << "this is a common engine"; }

    SophonModel::~SophonModel() { UnInit(); }

    EVStatus SophonModel::UnInit()
    {
        if (m_init_flag)
        {
            EVLOG(INFO) << "uninit OvModel";
            m_net_ptr.reset();
            m_ctx_ptr.reset();
            m_net_names.clear();
            m_handle_ptr.reset();
            m_init_flag = false;
            return EV_SUCCESS;
        }
        return EV_REPEATED_OP;
    }

    // init model and infos
    EVStatus SophonModel::Init(const char *net, void **config)
    {
        if (m_init_flag)
        {
            EVLOG(WARNING) << "call init mutiple times";
            return EV_REPEATED_OP;
        }

        ToolReport report;
        // report info
        int code = report.report(3, "ev_model_common-arm-sophon", "V1.0.0", "common model interface, sophon version");
        if (code == 0)
        {
            EVLOG(INFO) << "report done successully";
        }
        std::string model_path = std::string(net);
        if (model_path.find(".bmodel") == std::string::npos)
        {
            model_path = model_path.substr(0, model_path.find_last_of('.')) + ".bmodel";
        }
        if (CheckFileType(model_path.c_str()) != EV_FILE)
        {
            EVLOG(ERROR) << "invalid model or not exist:" << model_path; // 只能加载.bmodel
            exit(-1);
        }
        LoadBmModel(std::string(model_path));
        if (m_is_ok)
        {
            m_init_flag = true;
            for (size_t i = 0; i < m_net_ptr->Inputs().size(); ++i)
            {
                const auto &input_tensor = m_net_ptr->Inputs()[i]->tensor();
                std::string dims_info = "[";
                for (int j = 0; j < input_tensor->shape.num_dims; ++j)
                {
                    dims_info += std::to_string(input_tensor->shape.dims[j]);
                    if (j != input_tensor->shape.num_dims - 1)
                    {
                        dims_info += ",";
                    }
                    else
                    {
                        dims_info += "]";
                    }
                }
                EVLOG(INFO) << "input " << i << " shape:" << dims_info;
            }
            auto output_layer_size = m_net_ptr->Outputs().size();
            m_data_outputs.resize(output_layer_size);
            for (size_t i = 0; i < m_net_ptr->Outputs().size(); ++i)
            {
                const auto &output_tensor = m_net_ptr->Outputs()[i]->tensor();
                size_t elem_size{1};
                std::string dims_info = "[";
                for (int j = 0; j < output_tensor->shape.num_dims; ++j)
                {
                    elem_size *= output_tensor->shape.dims[j];
                    dims_info += std::to_string(output_tensor->shape.dims[j]);
                    if (j != output_tensor->shape.num_dims - 1)
                    {
                        dims_info += ",";
                    }
                    else
                    {
                        dims_info += "]";
                    }
                }
                m_data_outputs[i].resize(elem_size);
                EVLOG(INFO) << "output " << i << " shape:" << dims_info;
            }
        }
        return EV_SUCCESS;
    }

    // run inference
    EVStatus SophonModel::Infer(EVModelData *in, EVModelData *out)
    {
        if (!m_init_flag)
        {
            EVLOG(ERROR) << "infer failed, model was not initialized";
            return EV_FAIL;
        }
        REC_TIME(t0);
        if (static_cast<size_t>(in->mat_num) != m_net_ptr->Inputs().size())
        {
            EVLOG(ERROR) << "Model requires: " << m_net_ptr->Inputs().size() << " inputs but " << in->mat_num
                         << " is provided.";
            return EV_FAIL;
        }
        for (size_t i = 0; i < m_net_ptr->Inputs().size(); i++)
        {
            if (static_cast<size_t>(in->mat[i].data_size) != m_net_ptr->Inputs()[i]->ByteSize())
            {
                EVLOG(ERROR) << "input size mismatch : " << i << "," << in->mat[i].data_size << " vs "
                             << m_net_ptr->Inputs()[i]->ByteSize();
                return EV_FAIL;
            }
            m_net_ptr->Inputs()[i]->Reshape(m_net_ptr->Inputs()[i]->tensor()->shape);
            auto ret = m_net_ptr->Inputs()[i]->CopyFrom(in->mat[i].data);
            if (ret != BM_SUCCESS)
            {
                EVLOG(ERROR) << "CopyFrom error";
                return EV_FAIL;
            }
        }
        m_net_ptr->Forward();
        auto output_layer_size = m_net_ptr->Outputs().size();
        out->mat_num = output_layer_size;
        out->mat = new EVMatData[output_layer_size];
        for (size_t i = 0; i < m_net_ptr->Outputs().size(); ++i)
        {
            const auto &output_tensor = m_net_ptr->Outputs()[i]->tensor();
            size_t elem_size{1};
            for (int j = 0; j < output_tensor->shape.num_dims; ++j)
            {
                elem_size *= output_tensor->shape.dims[j];
                out->mat[i].dims.push_back(output_tensor->shape.dims[j]);
            }
            m_net_ptr->Outputs()[i]->CopyTo(m_data_outputs[i].data());
            out->mat[i].type = EVDataFormat::EV_FLOAT;
            out->mat[i].loc = EV_DATA_HOST;
            out->mat[i].data_size = elem_size * sizeof(float);
            out->mat[i].data = m_data_outputs[i].data();
            out->mat[i].desc = m_output_names[i];
        }
        REC_TIME(t1);
        EVLOG(INFO) << "inference time(ms):" << RUN_TIME(t1 - t0);
        return EV_SUCCESS;
    }

    EVStatus SophonModel::OptimizeModel(const char *model_name) { return EV_SUCCESS; }

    EVStatus SophonModel::SetIO(const std::vector<std::string> &inputs, const std::vector<std::string> &outputs)
    {
        m_input_names = inputs;
        m_output_names = outputs;
        EVLOG(INFO) << "set IO done";
        return EV_SUCCESS;
    }

    CREATE_INFER_INS(SophonModel)
    DESTROY_INFER_INS(SophonModel)

} // namespace ev
